/*
	Yelo: Open Sauce SDK
		Halo 1 (CE) Edition

	See license\OpenSauce\Halo1_CE for specific license information
*/

namespace Yelo
{
	namespace Objects
	{
		namespace UnitUpgrades
		{
			// Checks if the given seat index has multiple boarding seats and if anyone is in them
			// Allows us to prevent boarding when someone is already boarding the given vehicle's seat_index
			static bool GetUnitInBoardingSeats(datum_index vehicle_index, int32 seat_index)
			{
				bool enterable = true;
				s_unit_datum* vehicle_unit = (*Objects::ObjectHeader())[vehicle_index]->Type._unit;
				
				datum_index vehicle_tag_index = *vehicle_unit->object.GetTagDefinition();
				int32 unit_upgrade_definition_index = DefinitionFindUpgradableUnitIndex(vehicle_tag_index);

				// If unit upgrades exist for the given vehicle, check if anyone else is boarding
				if (unit_upgrade_definition_index != NONE)
				{
					TagGroups::s_unit_upgrades_definition const& unit_upgrade = 
						TagGroups::_global_yelo_globals->unit_upgrades[unit_upgrade_definition_index];

					// Check each seat of our vehicle and see if there's someone in one of the other boarding seats
					for (int x = 0; x < unit_upgrade.seats.Count; x++)
					{
						int16 seat_boarded = unit_upgrade.seats[x].seat_boarded;

						// If the seat_boarded == seat_index, check if anyone is in this seat
						if (seat_boarded == seat_index)
						{
							datum_index unit_index = GetUnitInSeat(vehicle_index, x);

							if (unit_index != datum_index::null)
								enterable = false;
						}
					}
				}

				return enterable;
			}

			// Resets the unit flags we set during boarding
			static void UnitResetBoardingControlFlags(s_unit_datum* unit)
			{
				datum_index vehicle_index = *unit->object.GetParentObjectIndex();
				s_unit_datum* vehicle = (*Objects::ObjectHeader())[vehicle_index]->Type._unit;
				
				SET_FLAG(*unit->unit.GetFlags(), 0, true);
				SET_FLAG(*vehicle->unit.GetFlags(), 16, false);
			}

			// Compares the seat blocks in project_yellow_globals and the vehicle tag to make sure they are the same
			static bool CheckUnitUpgradeSeats(TagGroups::s_unit_upgrades_definition const& unit_upgrade, 
				datum_index unit_index)
			{
				const TagGroups::s_unit_definition* unit_definition = GetObjectDefinition<TagGroups::s_unit_definition>(unit_index);

				// If the seat count matches, compare seat block fields
				if (unit_definition->unit.seats.Count == unit_upgrade.seats.Count)
				{
					// Compare the seat labels, marker names, and flags of the unit_definition and unit_upgrade blocks
					for (int x = 0; x < unit_upgrade.seats.Count; x++)
					{
						if (strcmp(unit_definition->unit.seats[x].label, unit_upgrade.seats[x].label) != 0)
							return false;
						if (strcmp(unit_definition->unit.seats[x].marker_name, unit_upgrade.seats[x].marker_name) != 0)
							return false;
						if (unit_definition->unit.seats[x].flags.driver_bit != unit_upgrade.seats[x].flags.driver_bit)
							return false;
					}
				}
				else
					return false;

				return true;
			}

			// Starts the boarding sequence on the boarding unit and kicks the boarded unit from the boarded_seat_index
			static void BoardingStart(TagGroups::s_unit_upgrades_definition const& unit_upgrade, 
				datum_index vehicle_index, datum_index boarding_unit_index, datum_index boarded_unit_index)
			{
				TagGroups::model_animation_graph const* boarding_animation_graph = GetObjectAnimations(boarding_unit_index);
				s_unit_datum* boarding_unit = (*Objects::ObjectHeader())[boarding_unit_index]->Type._unit;
				s_unit_datum* vehicle = (*Objects::ObjectHeader())[vehicle_index]->Type._unit;
				int16 boarding_seat_index = *boarding_unit->unit.GetVehicleSeatIndex();
				int16 boarded_seat_index = unit_upgrade.seats[boarding_seat_index].seat_boarded;

				// We use the primary keyframe index to trigger the eject sequence of the boarded unit
				datum_index boarding_unit_animation_definition = *boarding_unit->object.GetAnimationDefinition();
				int32 enter_animation_index = UnitGetAnimationIndexFromSeatAnimation(boarding_unit_index, 
					Enums::unit_seat_animation::_unit_seat_animation_enter);
				int32 board_animation_index = UnitGetAnimationIndexFromSeatAnimation(boarding_unit_index, 
					Enums::unit_seat_animation::_unit_seat_animation_yelo_board);
				int16 primary_keyframe_index = boarding_animation_graph->animations[board_animation_index].key_frame_index;
				int16 boarding_unit_animation_frames_left = ObjectGetAnimationFramesLeft(boarding_unit_index);
				int16* boarding_unit_animation_index = boarding_unit->object.GetAnimationCurrentIndex();
				int16* boarding_unit_animation_frame_index = boarding_unit->object.GetAnimationCurrentFrameIndex();

				// If the boarding_unit is killed during boarding, 
				// enable player entry into the vehicle
				if (TEST_FLAG(*boarding_unit->object.GetDamageFlags(), 5))
					SET_FLAG(*vehicle->unit.GetFlags(), 16, false);
				// Else, disable player entry into the vehicle
				else
					SET_FLAG(*vehicle->unit.GetFlags(), 16, true);

				// Enter boarding seat
				if (*boarding_unit_animation_index == enter_animation_index && 
					*boarding_unit_animation_frame_index == 0)
				{
					// Reset replacement animations during boarding
					*boarding_unit->unit.GetAnimation()->GetReplacementAnimationState() = 0;
					*boarding_unit->unit.GetAnimation()->GetReplacementAnimationIndex() = NONE;
					*boarding_unit->unit.GetAnimation()->GetReplacementAnimationFrameIndex() = 0;
					// Reset overlay animations during boarding
					*boarding_unit->unit.GetAnimation()->GetOverlayAnimationState() = 0;
					*boarding_unit->unit.GetAnimation()->GetOverlayAnimationIndex() = NONE;
					*boarding_unit->unit.GetAnimation()->GetOverlayAnimationFrameIndex() = 0;
					// Disable player actions while boarding
					SET_FLAG(*boarding_unit->unit.GetFlags(), 0, false);
				}
				// Play the boarding animation
				else if (*boarding_unit_animation_index == enter_animation_index && 
						 boarding_unit_animation_frames_left == 0)
				{
					// Interpolate animations
					Engine::Objects::StartInterpolation(boarding_unit_index, 6);
					// Open the vehicle
					Engine::Objects::UnitOpen(vehicle_index);
					// Play ejection animation
					Engine::Objects::UnitSetAnimation(boarding_unit_index, 
						boarding_unit_animation_definition, board_animation_index);
				}
				// Check for a boarded_unit
				else if (boarded_unit_index != datum_index::null)
				{
					s_unit_datum* boarded_unit = (*Objects::ObjectHeader())[boarded_unit_index]->Type._unit;
					datum_index boarded_unit_animation_definition = *boarded_unit->object.GetAnimationDefinition();
					sbyte* boarded_unit_animation_state = boarded_unit->unit.GetAnimation()->GetUnitAnimationState();
					int16* boarded_unit_animation_index = boarded_unit->object.GetAnimationCurrentIndex();
					int16* boarded_unit_animation_frame_index = boarded_unit->object.GetAnimationCurrentFrameIndex();
					int32 eject_animation_index = UnitGetAnimationIndexFromSeatAnimation(boarded_unit_index, 
						Enums::unit_seat_animation::_unit_seat_animation_yelo_eject);
				
					// Eject the boarded_unit on the primary keyframe index of the boarding animation
					if (*boarded_unit_animation_index != eject_animation_index && 
						*boarding_unit_animation_index == board_animation_index &&
						*boarding_unit_animation_frame_index >= primary_keyframe_index)
					{
						// Play Ejection Animation and exit the vehicle
						Engine::Objects::UnitSetAnimation(boarded_unit_index, 
							boarded_unit_animation_definition, eject_animation_index);
						*boarded_unit_animation_state = 
							Enums::unit_animation_state::_unit_animation_state_seat_exit;
						
						// If boarding is set to kill the passenger, kill them
						if (unit_upgrade.seats[boarding_seat_index].flags.boarding_kills_passenger_bit)
							Engine::Objects::DepleteBody(boarded_unit_index);
						// If boarding is set to deplete shields, deplete them
						if (unit_upgrade.seats[boarding_seat_index].flags.boarding_depletes_shield_bit)
							Engine::Objects::DepleteShield(boarded_unit_index);
					}
				}
			}

			// Ends the boarding sequence and forces the boarding unit into the seat_boarded if boarding is set to enter the seat
			static void BoardingEnd(TagGroups::s_unit_upgrades_definition const& unit_upgrade, 
				datum_index vehicle_index, datum_index boarding_unit_index)
			{
				s_unit_datum* boarding_unit = (*Objects::ObjectHeader())[boarding_unit_index]->Type._unit;
				s_unit_datum* vehicle = (*Objects::ObjectHeader())[vehicle_index]->Type._unit;
				datum_index boarding_unit_animation_definition = *boarding_unit->object.GetAnimationDefinition();
				sbyte* boarding_unit_animation_state = boarding_unit->unit.GetAnimation()->GetUnitAnimationState();
				int16* boarding_unit_animation_index = boarding_unit->object.GetAnimationCurrentIndex();
				int16* boarding_unit_animation_frame_index = boarding_unit->object.GetAnimationCurrentFrameIndex();
				int16 boarding_unit_animation_frames_left = ObjectGetAnimationFramesLeft(boarding_unit_index);
				int16 boarding_seat_index = *boarding_unit->unit.GetVehicleSeatIndex();
				int16 boarded_seat_index = unit_upgrade.seats[boarding_seat_index].seat_boarded;
				int32 board_animation_index = UnitGetAnimationIndexFromSeatAnimation(boarding_unit_index, 
					Enums::unit_seat_animation::_unit_seat_animation_yelo_board);

				// If the boarded_seat_index is a drivers seat and we're entering it, keep the engine running
				if (unit_upgrade.seats[boarding_seat_index].flags.boarding_enters_seat_bit && 
					unit_upgrade.seats[boarded_seat_index].flags.driver_bit)
				{
					*vehicle->unit.GetDriverObj() = boarding_unit_index;
					*vehicle->unit.GetDriverPower() = 1.0f;
				}

				// If the board_animation is complete
				if (*boarding_unit_animation_index == board_animation_index && 
					boarding_unit_animation_frames_left == 0)
				{
					// Make the boarding_unit exit the vehicle
					*boarding_unit_animation_state = Enums::unit_animation_state::_unit_animation_state_seat_exit;

					// If boarding enters the seat put the boarding_unit in seat_boarded
					if (unit_upgrade.seats[boarding_seat_index].flags.boarding_enters_seat_bit && 
						unit_upgrade.seats[boarding_seat_index].seat_boarded != NONE)
					{
						Engine::Objects::UnitEnterVehicleSeatIndex(boarding_unit_index, vehicle_index, boarded_seat_index);
							
						// Begin the idle animation after we enter for a smooth transition into the seat
						int32 idle_animation_index = UnitGetAnimationIndexFromWeaponClass(boarding_unit_index, 
							Enums::weapon_class_animation::_weapon_class_animation_idle);
						Engine::Objects::UnitSetAnimation(boarding_unit_index, 
							boarding_unit_animation_definition, idle_animation_index);
						*boarding_unit_animation_state = 
							Enums::unit_animation_state::_unit_animation_state_custom_animation;
					}

					// Close the vehicle
					Engine::Objects::UnitClose(vehicle_index);
					// Enable player actions
					SET_FLAG(*boarding_unit->unit.GetFlags(), 0, true);
					// Enable player entry into the vehicle
					SET_FLAG(*vehicle->unit.GetFlags(), 16, false);
				}
			}

			// Checks if the vehicle has unit upgrades and if the seat the entering_unit is in (or is trying to enter) 
			// is a boarding seat and determines if we should start the boarding sequence
			static void BoardingCheckSeat(datum_index entering_unit_index)
			{
				s_unit_datum* entering_unit = (*Objects::ObjectHeader())[entering_unit_index]->Type._unit;
				datum_index vehicle_index = *entering_unit->object.GetParentObjectIndex();
				
				if (TagGroups::_global_yelo_globals->unit_upgrades.Count == 0 || 
					entering_unit_index == datum_index::null || 
					vehicle_index == datum_index::null)
					return;
				
				// Find the s_unit_upgrade definition for the vehicle we're currently entering
				s_unit_datum* vehicle = (*Objects::ObjectHeader())[vehicle_index]->Type._unit;
				datum_index vehicle_tag_index = *vehicle->object.GetTagDefinition();
				int32 unit_upgrade_definition_index = DefinitionFindUpgradableUnitIndex(vehicle_tag_index);

				// If there are unit_upgrades defined for the vehicle we're entering, check for boarding seats
				if (unit_upgrade_definition_index != NONE)
				{
					int16 current_seat_index = *entering_unit->unit.GetVehicleSeatIndex();
					TagGroups::s_unit_upgrades_definition const& unit_upgrade = 
						TagGroups::_global_yelo_globals->unit_upgrades[unit_upgrade_definition_index];

					// Make sure the seat blocks in the vehicle tag and unit_upgrade block are exactly the same
					if (CheckUnitUpgradeSeats(unit_upgrade, vehicle_index) == false)
						return;

					// If the current seat is a boarding seat, get the boarding/boarded unit datums
					if (current_seat_index != NONE && unit_upgrade.seats[current_seat_index].flags.boarding_seat_bit)
					{
						int16 boarding_seat_index = current_seat_index;
						int16 boarded_seat_index = unit_upgrade.seats[boarding_seat_index].seat_boarded;
						datum_index boarding_unit_index = GetUnitInSeat(vehicle_index, boarding_seat_index);
						datum_index boarded_unit_index = GetUnitInSeat(vehicle_index, boarded_seat_index);
						
						if (boarding_unit_index != datum_index::null)
						{
							// Start the boarding sequence
							BoardingStart(unit_upgrade, vehicle_index, boarding_unit_index, boarded_unit_index);

							// If boarded_unit_index is null, they have been successfully 
							// kicked from the vehicle, end the boarding sequence
							if (boarded_unit_index == datum_index::null)
								BoardingEnd(unit_upgrade, vehicle_index, boarding_unit_index);
						}
					}
				}
			}

			// Overwrites the engine's UnitCanEnterSeat check to determine if we can enter the given seat index
			// multiteam vehicles will always be enabled using this
			static bool UnitCanEnterSeatCheck(datum_index entering_unit_index, datum_index vehicle_index, int16 seat_index, bool enterable)
			{
				if (entering_unit_index != datum_index::null && vehicle_index != datum_index::null)
				{
					s_unit_datum* entering_unit = (*Objects::ObjectHeader())[entering_unit_index]->Type._unit;
					s_unit_datum* vehicle_unit = (*Objects::ObjectHeader())[vehicle_index]->Type._unit;
					datum_index seated_unit = GetUnitInSeat(vehicle_index, seat_index);

					datum_index vehicle_tag_index = *vehicle_unit->object.GetTagDefinition();
					int32 unit_upgrade_definition_index = DefinitionFindUpgradableUnitIndex(vehicle_tag_index);

					// Check to see if the vehicle has upgrades
					if (unit_upgrade_definition_index != NONE)
					{
						TagGroups::s_unit_upgrades_definition const& unit_upgrade = 
							TagGroups::_global_yelo_globals->unit_upgrades[unit_upgrade_definition_index];

						// If the seat we're trying to enter is a boarding seat
						if (unit_upgrade.seats[seat_index].flags.boarding_seat_bit)
						{
							int16 boarding_seat_index = seat_index;
							int16 boarded_seat_index = unit_upgrade.seats[boarding_seat_index].seat_boarded;
							datum_index boarding_unit_index = entering_unit_index;
							datum_index boarded_unit_index = GetUnitInSeat(vehicle_index, boarded_seat_index);
							
							// Check to see if anyone else is boarding the vehicle
							if (!GetUnitInBoardingSeats(vehicle_index, boarded_seat_index))
								enterable = false;

							// If there is a unit to be boarded we can enter the boarding seat
							if (boarded_unit_index != datum_index::null)
							{
								s_unit_datum* boarding_unit = (*Objects::ObjectHeader())[boarding_unit_index]->Type._unit;
								s_unit_datum* boarded_unit = (*Objects::ObjectHeader())[boarded_unit_index]->Type._unit;
								int16 boarding_team = *boarding_unit->object.GetOwnerTeamIndex();
								int16 boarded_team = *boarded_unit->object.GetOwnerTeamIndex();

								// If the boarded_unit is on our team, we can't enter the boarding seat
								if (!Engine::Game::TeamIsEnemy(boarded_team, boarding_team))
									enterable = false;
							}
							// If there is no unit to be boarded we can't enter the boarding seat
							else
								enterable = false;
						}
					}

					// Check if someone's in the seat we're trying to enter
					if (seated_unit != datum_index::null)
						enterable = false;
				}

				return enterable;
			}

			// Hooks the switch case for the seat_exit animation state
			// Allows us to reset our boarding control flags when a player exits a vehicle
			static API_FUNC_NAKED void PLATFORM_API UnitSeatExitAnimationStateHook()
			{
				static uint32 RETN_ADDRESS = GET_FUNC_PTR(UNIT_SEAT_EXIT_ANIMATION_STATE_RETN);

				__asm {
					push	eax
					push	ecx

					push	ebx			// s_unit_datum* exiting_unit
					call	UnitResetBoardingControlFlags

					pop		ecx
					pop		eax

					mov     edx, [ebx+11Ch]

					jmp		RETN_ADDRESS
				}
			}
			
			// Hooks the function that ejects a player from a flipped vehicle
			// Allows us to reset our boarding control flags when a player's vehicle flips and they are ejected
			static API_FUNC_NAKED void PLATFORM_API UnitSeatEjectionHook()
			{
				static uint32 RETN_ADDRESS = GET_FUNC_PTR(UNIT_SEAT_EJECTION_RETN);

				__asm {
					push	edx

					push	ebx			// s_unit_datum* ejected_unit
					call	UnitResetBoardingControlFlags

					pop		edx

					mov     eax, [ebx+11Ch]

					jmp		RETN_ADDRESS
				}
			}

			// Called when a unit is in a vehicle
			static API_FUNC_NAKED void PLATFORM_API UnitBoardingCheckSeatHook()
			{
				static uint32 RETN_ADDRESS = GET_FUNC_PTR(UNIT_BOARDING_CHECK_SEAT_RETN);

				__asm {
					push	eax

					push	esi		// datum_index entering_unit_index
					call	BoardingCheckSeat

					pop		eax
					test    byte ptr [ebx+208h], 40h

					jmp		RETN_ADDRESS
				}
			}

			// Overwrites the engines UnitCanEnterSeat function with our own to compensate for boarding rules
			static API_FUNC_NAKED void PLATFORM_API UnitCanEnterSeatHook()
			{
				static uint32 RETN_ADDRESS = GET_FUNC_PTR(UNIT_CAN_ENTER_SEAT_RETN);

				__asm {
					xor		eax, eax
					movsx	eax, bl

					push	eax				// bool enterable
					mov		eax, [esp+28]
					push	eax				// int16 seat_index
					push	edx				// datum_index vehicle_index
					push	esi				// datum_index entering_unit_index
					call	UnitCanEnterSeatCheck
					mov		ebx, eax		// mov the resulting bool to ebx
					
					mov		eax, [esp-8]
					mov		edx, [esp-12]
					push	eax				// int16 seat_index
					push	edx				// datum_index vehicle_index
					call	GetUnitInSeat	// In the original code, if a unit is in the seat
											// we're trying to enter, the unit's datum_index is
											// saved to ebp

					mov		ebp, eax		// mov the resulting datum_index to ebp
					jmp		RETN_ADDRESS
				}
			}

			void Initialize()
			{
				Memory::WriteRelativeJmp(&UnitSeatExitAnimationStateHook, GET_FUNC_VPTR(UNIT_SEAT_EXIT_ANIMATION_STATE_HOOK), true);
				Memory::WriteRelativeJmp(&UnitSeatEjectionHook, GET_FUNC_VPTR(UNIT_SEAT_EJECTION_HOOK), true);
				Memory::WriteRelativeJmp(&UnitBoardingCheckSeatHook, GET_FUNC_VPTR(UNIT_BOARDING_CHECK_SEAT_HOOK), true);
				Memory::WriteRelativeJmp(&UnitCanEnterSeatHook, GET_FUNC_VPTR(UNIT_CAN_ENTER_SEAT_HOOK), true); 

				// Prevents AI from exiting a vehicle while being boarded
				AI::AIExitEnemyVehicleFix();
			}

			void Dispose()
			{
			}
		};
	};
};